import numpy as np
import rospy
from unitree_legged_msgs.msg import LowState, MotorState
from sensor_msgs.msg import Imu
from geometry_msgs.msg import WrenchStamped


class multiThread():
    servo_sub = list(np.arange(12))
    footForce_sub = list(np.arange(4))

    def __init__(self, rname, lowState):
        self.robot_name = rname
        self.lowState = lowState
        self.imu_sub = rospy.Subscriber("/trunk_imu", Imu, self.imuCallback)
        self.footForce_sub[0] = rospy.Subscriber("/visual/FR_foot_contact/the_force", WrenchStamped, self.FRfootCallback)
        self.footForce_sub[1] = rospy.Subscriber("/visual/FL_foot_contact/the_force", WrenchStamped, self.FLfootCallback)
        self.footForce_sub[2] = rospy.Subscriber("/visual/RR_foot_contact/the_force", WrenchStamped, self.RRfootCallback)
        self.footForce_sub[3] = rospy.Subscriber("/visual/RL_foot_contact/the_force", WrenchStamped, self.RLfootCallback)
        self.servo_sub[0] = rospy.Subscriber("/" + self.robot_name + "_gazebo/FR_hip_controller/state", MotorState, self.FRhipCallback)
        self.servo_sub[1] = rospy.Subscriber("/" + self.robot_name + "_gazebo/FR_thigh_controller/state", MotorState, self.FRthighCallback)
        self.servo_sub[2] = rospy.Subscriber("/" + self.robot_name + "_gazebo/FR_calf_controller/state", MotorState, self.FRcalfCallback)
        self.servo_sub[3] = rospy.Subscriber("/" + self.robot_name + "_gazebo/FL_hip_controller/state", MotorState, self.FLhipCallback)
        self.servo_sub[4] = rospy.Subscriber("/" + self.robot_name + "_gazebo/FL_thigh_controller/state", MotorState, self.FLthighCallback)
        self.servo_sub[5] = rospy.Subscriber("/" + self.robot_name + "_gazebo/FL_calf_controller/state", MotorState, self.FLcalfCallback)
        self.servo_sub[6] = rospy.Subscriber("/" + self.robot_name + "_gazebo/RR_hip_controller/state", MotorState, self.RRhipCallback)
        self.servo_sub[7] = rospy.Subscriber("/" + self.robot_name + "_gazebo/RR_thigh_controller/state", MotorState, self.RRthighCallback)
        self.servo_sub[8] = rospy.Subscriber("/" + self.robot_name + "_gazebo/RR_calf_controller/state", MotorState, self.RRcalfCallback)
        self.servo_sub[9] = rospy.Subscriber("/" + self.robot_name + "_gazebo/RL_hip_controller/state", MotorState, self.RLhipCallback)
        self.servo_sub[10] = rospy.Subscriber("/" + self.robot_name + "_gazebo/RL_thigh_controller/state", MotorState, self.RLthighCallback)
        self.servo_sub[11] = rospy.Subscriber("/" + self.robot_name + "_gazebo/RL_calf_controller/state", MotorState, self.RLcalfCallback)
    
    def imuCallback(self, msg:Imu):
        self.lowState.imu.quaternion[0] = msg.orientation.w
        self.lowState.imu.quaternion[1] = msg.orientation.x
        self.lowState.imu.quaternion[2] = msg.orientation.y
        self.lowState.imu.quaternion[3] = msg.orientation.z

        self.lowState.imu.gyroscope[0] = msg.angular_velocity.x
        self.lowState.imu.gyroscope[1] = msg.angular_velocity.y
        self.lowState.imu.gyroscope[2] = msg.angular_velocity.z
        
        self.lowState.imu.accelerometer[0] = msg.linear_acceleration.x
        self.lowState.imu.accelerometer[1] = msg.linear_acceleration.y
        self.lowState.imu.accelerometer[2] = msg.linear_acceleration.z


    def FRhipCallback(self, msg:MotorState):
        self.start_up = False
        self.lowState.motorState[0].mode = msg.mode
        self.lowState.motorState[0].q = msg.q                # current angle (unit: radian)
        self.lowState.motorState[0].dq = msg.dq              # current velocity (unit: radian/second)
        self.lowState.motorState[0].tauEst = msg.tauEst      # current estimated output torque (unit: N.m)

    def FRthighCallback(self, msg:MotorState):
        self.lowState.motorState[1].mode = msg.mode
        self.lowState.motorState[1].q = msg.q
        self.lowState.motorState[1].dq = msg.dq
        self.lowState.motorState[1].tauEst = msg.tauEst

    def FRcalfCallback(self, msg:MotorState):
        self.lowState.motorState[2].mode = msg.mode
        self.lowState.motorState[2].q = msg.q
        self.lowState.motorState[2].dq = msg.dq
        self.lowState.motorState[2].tauEst = msg.tauEst

    def FLhipCallback(self, msg:MotorState):
        self.start_up = False
        self.lowState.motorState[3].mode = msg.mode
        self.lowState.motorState[3].q = msg.q
        self.lowState.motorState[3].dq = msg.dq
        self.lowState.motorState[3].tauEst = msg.tauEst

    def FLthighCallback(self, msg:MotorState):
        self.lowState.motorState[4].mode = msg.mode
        self.lowState.motorState[4].q = msg.q
        self.lowState.motorState[4].dq = msg.dq
        self.lowState.motorState[4].tauEst = msg.tauEst

    def FLcalfCallback(self, msg:MotorState):
        self.lowState.motorState[5].mode = msg.mode
        self.lowState.motorState[5].q = msg.q
        self.lowState.motorState[5].dq = msg.dq
        self.lowState.motorState[5].tauEst = msg.tauEst

    def RRhipCallback(self, msg:MotorState):
        self.start_up = False
        self.lowState.motorState[6].mode = msg.mode
        self.lowState.motorState[6].q = msg.q
        self.lowState.motorState[6].dq = msg.dq
        self.lowState.motorState[6].tauEst = msg.tauEst

    def RRthighCallback(self, msg:MotorState):
        self.lowState.motorState[7].mode = msg.mode
        self.lowState.motorState[7].q = msg.q
        self.lowState.motorState[7].dq = msg.dq
        self.lowState.motorState[7].tauEst = msg.tauEst

    def RRcalfCallback(self, msg:MotorState):
        self.lowState.motorState[8].mode = msg.mode
        self.lowState.motorState[8].q = msg.q
        self.lowState.motorState[8].dq = msg.dq
        self.lowState.motorState[8].tauEst = msg.tauEst

    def RLhipCallback(self, msg:MotorState):
        self.start_up = False
        self.lowState.motorState[9].mode = msg.mode
        self.lowState.motorState[9].q = msg.q
        self.lowState.motorState[9].dq = msg.dq
        self.lowState.motorState[9].tauEst = msg.tauEst

    def RLthighCallback(self, msg:MotorState):
        self.lowState.motorState[10].mode = msg.mode
        self.lowState.motorState[10].q = msg.q
        self.lowState.motorState[10].dq = msg.dq
        self.lowState.motorState[10].tauEst = msg.tauEst

    def RLcalfCallback(self, msg:MotorState):
        self.lowState.motorState[11].mode = msg.mode
        self.lowState.motorState[11].q = msg.q
        self.lowState.motorState[11].dq = msg.dq
        self.lowState.motorState[11].tauEst = msg.tauEst

    def FRfootCallback(self, msg:WrenchStamped):
        self.lowState.eeForce[0].x = msg.wrench.force.x
        self.lowState.eeForce[0].y = msg.wrench.force.y
        self.lowState.eeForce[0].z = msg.wrench.force.z
        self.lowState.footForce[0] = msg.wrench.force.z

    def FLfootCallback(self, msg:WrenchStamped):
        self.lowState.eeForce[1].x = msg.wrench.force.x
        self.lowState.eeForce[1].y = msg.wrench.force.y
        self.lowState.eeForce[1].z = msg.wrench.force.z
        self.lowState.footForce[1] = msg.wrench.force.z

    def RRfootCallback(self, msg:WrenchStamped):
        self.lowState.eeForce[2].x = msg.wrench.force.x
        self.lowState.eeForce[2].y = msg.wrench.force.y
        self.lowState.eeForce[2].z = msg.wrench.force.z
        self.lowState.footForce[2] = msg.wrench.force.z

    def RLfootCallback(self, msg:WrenchStamped):
        self.lowState.eeForce[3].x = msg.wrench.force.x
        self.lowState.eeForce[3].y = msg.wrench.force.y
        self.lowState.eeForce[3].z = msg.wrench.force.z
        self.lowState.footForce[3] = msg.wrench.force.z


class LowStateSubscriber():
    def __init__(self, rname):
        self.robot_name = rname
        self.lowState = LowState()
        # self.lowState_pub = rospy.Publisher("/" + self.robot_name + "_gazebo/lowState/state", LowState, queue_size=1)
        self.listen_publish_obj = multiThread(self.robot_name, self.lowState)
        # rospy.spin()
    
    def getState(self):
        return self.lowState
